import jittor as jt
from jittor import nn

from lib.Res2Net_v1b import res2net50_v1b_26w_4s


class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv(in_planes, out_planes, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm(out_planes)
        self.relu = nn.ReLU()

    def execute(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x


class RFB_modified(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(RFB_modified, self).__init__()
        self.relu = nn.ReLU()
        self.branch0 = nn.Sequential(BasicConv2d(in_channel, out_channel, 1))
        self.branch1 = nn.Sequential(BasicConv2d(in_channel, out_channel, 1), BasicConv2d(out_channel, out_channel, kernel_size=(1, 3), padding=(0, 1)), BasicConv2d(out_channel, out_channel, kernel_size=(3, 1), padding=(1, 0)), BasicConv2d(out_channel, out_channel, 3, padding=3, dilation=3))
        self.branch2 = nn.Sequential(BasicConv2d(in_channel, out_channel, 1), BasicConv2d(out_channel, out_channel, kernel_size=(1, 5), padding=(0, 2)), BasicConv2d(out_channel, out_channel, kernel_size=(5, 1), padding=(2, 0)), BasicConv2d(out_channel, out_channel, 3, padding=5, dilation=5))
        self.branch3 = nn.Sequential(BasicConv2d(in_channel, out_channel, 1), BasicConv2d(out_channel, out_channel, kernel_size=(1, 7), padding=(0, 3)), BasicConv2d(out_channel, out_channel, kernel_size=(7, 1), padding=(3, 0)), BasicConv2d(out_channel, out_channel, 3, padding=7, dilation=7))
        self.conv_cat = BasicConv2d((4 * out_channel), out_channel, 3, padding=1)
        self.conv_res = BasicConv2d(in_channel, out_channel, 1)

    def execute(self, x):
        x0 = self.branch0(x)
        x1 = self.branch1(x)
        x2 = self.branch2(x)
        x3 = self.branch3(x)
        x_cat = self.conv_cat(jt.contrib.concat((x0, x1, x2, x3), dim=1))
        x = nn.relu((x_cat + self.conv_res(x)))
        return x


class NeighborConnectionDecoder(nn.Module):
    def __init__(self, channel):
        super(NeighborConnectionDecoder, self).__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
        self.conv_upsample1 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample2 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample3 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample4 = BasicConv2d(channel, channel, 3, padding=1)
        self.conv_upsample5 = BasicConv2d((2 * channel), (2 * channel), 3, padding=1)
        self.conv_concat2 = BasicConv2d((2 * channel), (2 * channel), 3, padding=1)
        self.conv_concat3 = BasicConv2d((3 * channel), (3 * channel), 3, padding=1)
        self.conv4 = BasicConv2d((3 * channel), (3 * channel), 3, padding=1)
        self.conv5 = nn.Conv((3 * channel), 1, 1)

    def execute(self, x1, x2, x3):
        x1_1 = x1
        x2_1 = (self.conv_upsample1(self.upsample(x1)) * x2)
        x3_1 = ((self.conv_upsample2(self.upsample(x2_1)) * self.conv_upsample3(self.upsample(x2))) * x3)
        x2_2 = jt.contrib.concat((x2_1, self.conv_upsample4(self.upsample(x1_1))), dim=1)
        x2_2 = self.conv_concat2(x2_2)
        x3_2 = jt.contrib.concat((x3_1, self.conv_upsample5(self.upsample(x2_2))), dim=1)
        x3_2 = self.conv_concat3(x3_2)
        x = self.conv4(x3_2)
        x = self.conv5(x)
        return x


class GRA(nn.Module):
    def __init__(self, channel, subchannel):
        super(GRA, self).__init__()
        self.group = (channel // subchannel)
        self.conv = nn.Sequential(nn.Conv((channel + self.group), channel, 3, padding=1), nn.ReLU())
        self.score = nn.Conv(channel, 1, 3, padding=1)

    def execute(self, x, y):
        if (self.group == 1):
            x_cat = jt.contrib.concat((x, y), dim=1)
        elif (self.group == 2):
            xs = jt.chunk(x, 2, dim=1)
            x_cat = jt.contrib.concat((xs[0], y, xs[1], y), dim=1)
        elif (self.group == 4):
            xs = jt.chunk(x, 4, dim=1)
            x_cat = jt.contrib.concat((xs[0], y, xs[1], y, xs[2], y, xs[3], y), dim=1)
        elif (self.group == 8):
            xs = jt.chunk(x, 8, dim=1)
            x_cat = jt.contrib.concat((xs[0], y, xs[1], y, xs[2], y, xs[3], y, xs[4], y, xs[5], y, xs[6], y, xs[7], y), dim=1)
        elif (self.group == 16):
            xs = jt.chunk(x, 16, dim=1)
            x_cat = jt.contrib.concat((xs[0], y, xs[1], y, xs[2], y, xs[3], y, xs[4], y, xs[5], y, xs[6], y, xs[7], y, xs[8], y, xs[9], y, xs[10], y, xs[11], y, xs[12], y, xs[13], y, xs[14], y, xs[15], y), dim=1)
        elif (self.group == 32):
            xs = jt.chunk(x, 32, dim=1)
            x_cat = jt.contrib.concat((xs[0], y, xs[1], y, xs[2], y, xs[3], y, xs[4], y, xs[5], y, xs[6], y, xs[7], y, xs[8], y, xs[9], y, xs[10], y, xs[11], y, xs[12], y, xs[13], y, xs[14], y, xs[15], y, xs[16], y, xs[17], y, xs[18], y, xs[19], y, xs[20], y, xs[21], y, xs[22], y, xs[23], y, xs[24], y, xs[25], y, xs[26], y, xs[27], y, xs[28], y, xs[29], y, xs[30], y, xs[31], y), dim=1)
        else:
            raise Exception('Invalid Channel')
        x = (x + self.conv(x_cat))
        y = (y + self.score(x))
        return (x, y)


class ReverseStage(nn.Module):
    def __init__(self, channel):
        super(ReverseStage, self).__init__()
        self.weak_gra = GRA(channel, channel)
        self.medium_gra = GRA(channel, 8)
        self.strong_gra = GRA(channel, 1)

    def execute(self, x, y):
        y = (((- 1) * jt.sigmoid(y)) + 1)
        (x, y) = self.weak_gra(x, y)
        (x, y) = self.medium_gra(x, y)
        (_, y) = self.strong_gra(x, y)
        return y


class Network(nn.Module):
    def __init__(self, channel=32, imagenet_pretrained=False):
        super(Network, self).__init__()
        self.resnet = res2net50_v1b_26w_4s(pretrained=imagenet_pretrained)
        self.rfb2_1 = RFB_modified(512, channel)
        self.rfb3_1 = RFB_modified(1024, channel)
        self.rfb4_1 = RFB_modified(2048, channel)
        self.NCD = NeighborConnectionDecoder(channel)
        self.RS5 = ReverseStage(channel)
        self.RS4 = ReverseStage(channel)
        self.RS3 = ReverseStage(channel)

        self.upsample1_1 = nn.Upsample(scale_factor=8, mode='bilinear')
        self.upsample1_2 = nn.Upsample(scale_factor=0.25, mode='bilinear')

        self.upsample2_1 = nn.Upsample(scale_factor=32, mode='bilinear')
        self.upsample2_2 = nn.Upsample(scale_factor=2, mode='bilinear')

        self.upsample3_1 = nn.Upsample(scale_factor=16, mode='bilinear')
        self.upsample3_2 = nn.Upsample(scale_factor=2, mode='bilinear')

        self.upsample4 = nn.Upsample(scale_factor=8, mode='bilinear')

    def execute(self, x):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = nn.relu(x)
        x = self.resnet.maxpool(x)
        x1 = self.resnet.layer1(x)
        x2 = self.resnet.layer2(x1)
        x3 = self.resnet.layer3(x2)
        x4 = self.resnet.layer4(x3)
        x2_rfb = self.rfb2_1(x2)
        x3_rfb = self.rfb3_1(x3)
        x4_rfb = self.rfb4_1(x4)
        S_g = self.NCD(x4_rfb, x3_rfb, x2_rfb)
        S_g_pred = self.upsample1_1(S_g)
        guidance_g = self.upsample1_2(S_g)

        ra4_feat = self.RS5(x4_rfb, guidance_g)
        S_5 = (ra4_feat + guidance_g)

        S_5_pred = self.upsample2_1(S_5)
        guidance_5 = self.upsample2_2(S_5)

        ra3_feat = self.RS4(x3_rfb, guidance_5)
        S_4 = (ra3_feat + guidance_5)

        S_4_pred = self.upsample3_1(S_4)
        guidance_4 = self.upsample3_2(S_4)

        ra2_feat = self.RS3(x2_rfb, guidance_4)
        S_3 = (ra2_feat + guidance_4)

        S_3_pred = self.upsample4(S_3)

        return (S_g_pred, S_5_pred, S_4_pred, S_3_pred)


if (__name__ == '__main__'):
    import numpy as np
    from time import time
    net = Network(imagenet_pretrained=False)
    net.eval()

    dump_x = jt.randn(1, 3, 352, 352)
    frame_rate = np.zeros((1000, 1))
    for i in range(1000):
        start = time()
        y = net(dump_x)
        end = time()
        running_frame_rate = (1 * float((1 / (end - start))))
        print(i, '->', running_frame_rate)
        frame_rate[i] = running_frame_rate
    print(np.mean(frame_rate))
    print(y.shape)
